import readlineSync from "readline-sync";
import { ChatOllama } from "@langchain/ollama";
import {
  ChatPromptTemplate,
  HumanMessagePromptTemplate,
  SystemMessagePromptTemplate,
} from "@langchain/core/prompts";
import { ChatMessageHistory } from "langchain/stores/message/in_memory";
import {
  RunnableSequence,
  RunnableLambda,
  RunnablePassthrough,
} from "@langchain/core/runnables";
import { HumanMessage, AIMessage } from "@langchain/core/messages";
import { summaryChain } from "./summaryChain.js"; // 导入摘要链

// 1. 模型
const chatModel = new ChatOllama({ model: "llama3", temperature: 0.4 });

// 2. 提示词
const chatPrompt = ChatPromptTemplate.fromMessages([
  SystemMessagePromptTemplate.fromTemplate(`
  你是一名中文对话助手。无论用户使用何种语言，**必须始终用【简体中文】回答**。
  在作答前，先阅读下面可能提供的“长期摘要”，保证表述与历史一致、信息不自相矛盾。
  {maybe_summary_section}
  
  严格遵守：
  - 只用简体中文回答；除代码块外，禁止出现英文句子或段落。
  - 直接给出答案与步骤，避免客套和重复提示。
  - 列点请用精炼中文要点；必要时给出简短结论→理由→建议的结构。
  - 若需给代码，**代码内注释与说明也用中文**（保留必要的英文关键字/标识符即可）。
  - 问题含糊时，依据摘要与本轮上下文做最合理假设并继续回答，**不要反问**。
  
  示例（用于风格约束）：
  - ✅ “可以。建议你这样做：… 然后… 最后…”
  - ✅ “原因是：…；因此建议：…”
  - ❌ “Sure, here is …”
  - ❌ “Of course! …”
  `),
  HumanMessagePromptTemplate.fromTemplate("{input}"),
]);

let summary = ""; // 存储摘要信息的
const history = new ChatMessageHistory(); // 临时存储本轮用户对话

// 3. 构建对话链
const chatChain = RunnableSequence.from([
  // 记录用户的输入
  new RunnableLambda({
    func: async ({ input }) => {
      await history.addMessage(new HumanMessage(input));
      return { input };
    },
  }),
  // 注入摘要
  RunnablePassthrough.assign({
    maybe_summary_section: async () =>
      summary?.trim()
        ? `以下是之前对话的摘要，请在回答时保持一致性：\n${summary}\n`
        : "",
  }),
  chatPrompt,
  chatModel,
]);

async function chatLoop() {
  console.log("开始对话，输入内容后回车；/clear 清空摘要；/exit 退出。");

  while (true) {
    const input = readlineSync.question("\n你：").trim();
    if (!input) continue;

    // 退出功能
    if (input === "/exit" || input === "exit" || input === "quit") {
      console.log("已退出。");
      break;
    }

    // 清除历史
    if (input === "/clear") {
      summary = ""; // 清空摘要
      await history.clear(); // 清空 history
      console.log("已清空摘要与临时历史。");
      continue;
    }

    try {
      const stream = await chatChain.stream({ input });
      process.stdout.write("助理：");
      let finalText = "";
      for await (const chunk of stream) {
        process.stdout.write(chunk.content);
        finalText += chunk.content; // 拼接完整回复
      }
      console.log("\n");

      // 更新摘要
      // 将大模型返回的答案加入历史记录
      await history.addMessage(new AIMessage(finalText));

      // 只获取本轮的两条信息
      const lastTwo = (await history.getMessages()).slice(-2);
      // 通过本轮的消息对象构建一个本轮的消息字符串
      /**
       * `Human:你知道大象吗？
       * AI:我知道大象，大象是xxxx...`
       */
      const new_lines = lastTwo
        .map((m) => `${m.getType()}：${m.content}`)
        .join("\n");

      summary = await summaryChain.invoke({
        summary, // 之前的摘要
        new_lines, // 本轮的对话
      });
      await history.clear();

      // 打印新的摘要
      console.log("—— 当前摘要 ——");
      console.log(summary?.trim() ? summary.trim() : "（当前暂无摘要）");
      console.log("────────────────────────────");
    } catch (err) {
      console.error("调用模型失败：", err?.message ?? err);
    }
  }
}

chatLoop();
